
import struct
import ctypes
from ctypes import *
from typing import Union
from pathlib import Path

# 判断你的python环境是64位还是32位
__RuntimeEnvironment = struct.calcsize("P") * 8 == 64
import os

def find_file(file_name):
    for file_path in Path.cwd().rglob(file_name):
        return str(file_path)

dllpath = find_file("Sunny64.dll")
dllpath32 = find_file("Sunny.dll")




try:
    if __RuntimeEnvironment:
        # 如果是64位加载64位DLL
        lib = CDLL(dllpath)
        # Go语言回调函数声明
        TcpCallback = CFUNCTYPE(None, c_int64, c_char_p, c_char_p, c_int64, c_int64, c_int64, c_int64, c_int64, c_int64)
        HttpCallback = CFUNCTYPE(None, c_int64, c_int64, c_int64, c_int64, c_char_p, c_char_p, c_char_p, c_int64)
        WsCallback = CFUNCTYPE(None, c_int64, c_int64, c_int64, c_int64, c_char_p, c_char_p, c_int64, c_int64)
        UDPCallback = CFUNCTYPE(None, c_int64, c_char_p, c_char_p, c_int64, c_int64, c_int64, c_int64)
    else:
        # 如果不是64位加载32位DLL
        lib = CDLL(dllpath32)
        # Go语言回调函数声明
        TcpCallback = CFUNCTYPE(None, c_int, c_char_p, c_char_p, c_int, c_int, c_int, c_int, c_int, c_int)
        HttpCallback = CFUNCTYPE(None, c_int, c_int, c_int, c_int, c_char_p, c_char_p, c_char_p, c_int)
        WsCallback = CFUNCTYPE(None, c_int, c_int, c_int, c_int, c_char_p, c_char_p, c_int, c_int)
        UDPCallback = CFUNCTYPE(None, c_int, c_char_p, c_char_p, c_int, c_int, c_int, c_int)
except:
    print("载入DLL失败,请检测DLL文件")


# 这个类 是动态加载DLL时 设置返回值为指针
class LibSunny:
    def __getattr__(self, name):
        func = getattr(lib, name)
        func.restype = ctypes.POINTER(ctypes.c_int)
        return func


DLLSunny = LibSunny()


# 指针到字节数组 ptr=指针 skip=偏移数 num=取出几个字节
def PtrToByte(ptr, skip, num):
    result_as_int = ctypes.cast(ptr, ctypes.c_void_p).value
    if result_as_int == None:
        return bytearray()
    result_as_int += skip
    new_result_ptr = ctypes.cast(result_as_int, ctypes.POINTER(ctypes.c_int))
    buffer = ctypes.create_string_buffer(num)
    ctypes.memmove(buffer, new_result_ptr, num)
    return buffer.raw


# 指针到整数
def PtrToInt(ptr):
    return ctypes.cast(ptr, ctypes.c_void_p).value


# 指针到字符串
def PointerToText(ptr):
    if ptr == 0:
        return ""
    buff = b''
    i = 0
    while True:
        bs = PtrToByte(ptr, i, 1)
        i += 1
        if bs[0] == 0:
            break
        buff = buff + bs
    try:
        return buff.decode('utf-8')
    except:
        return buff.decode('gbk')


# 指针到字节数组 (DLL协商的前8个字节是长度)
def PointerToBytes(ptr):
    if ptr == 0:
        return bytearray()
    lp = PtrToByte(ptr, 0, 8)
    if len(lp) != 8:
        return lp
    Lxp = PtrToInt(DLLSunny.BytesToInt(create_string_buffer(lp), 8))
    return PtrToByte(ptr, 8, Lxp)


# http 请求操作类
class SunnyRequest:
    def __init__(self, _MessageId):
        self.MessageId = _MessageId

    def DeleteZipTag(self):
        '''
        删除压缩标记:请求协议头中去除Gzip 若不删除压缩标记，返回数据可能是压缩后的
        :return:
        '''
        DLLSunny.DelRequestHeader(self.MessageId, create_string_buffer("Accept-Encoding".encode("utf-8")))

    def setProxy(self, proxy):
        '''
        置代理：对这个请求使用指定代理请求请求  仅支持Socket5和http
        :param proxy: socket5://admin:123456@127.0.0.1:8888 或 http://admin:123456@127.0.0.1:8888
        :return:
        '''
        if not isinstance(proxy, str):
            return
        return bool(DLLSunny.SetRequestProxy(self.MessageId, create_string_buffer(proxy.encode("utf-8"))))

    def setOutTime(self, OutTime):
        '''
        置请求超时：仅限在发起请求时候使用
        :param OutTime: 单位【毫秒】
        :return:
        '''
        if not isinstance(OutTime, int):
            return
        DLLSunny.SetRequestProxy(self.MessageId, OutTime)

    def redirection(self, newUrl):
        '''
        请求重定向：可以转向  1->2   网址A->网址B
        :param newUrl: 欲转向地址
        :return:
        '''
        if not isinstance(newUrl, str):
            return
        DLLSunny.SetRequestUrl(self.MessageId, create_string_buffer(newUrl.encode("utf-8")))

    def modifyBody(self, newBody: Union[str, bytes]):
        '''
        修改Body数据
        :param newBody:可以是 str 或 bytes
        :return: 成功返回真，失败返回假
        '''
        if isinstance(newBody, str):
            newBody = newBody.encode()
        return PtrToInt(DLLSunny.SetRequestData(self.MessageId, create_string_buffer(newBody), len(newBody))) == 1

    def modifyHeaders(self, nweHeaders: Union[str, dict] = None, key: str = None, value: str = None):
        '''
        新增或修改协议头 值为空字符串时 删除该条协议头
        :param nweHeaders: nweHeaders: 可以是 str 或 dict 多条换行
        :param key: 单条 Key
        :param value: 单条 value
        :return: 成功返回真，失败返回假
        '''
        if nweHeaders:
            if isinstance(nweHeaders, str):
                my_list = nweHeaders.split("\r\n")
                for item in my_list:
                    my_list2 = item.split(":")
                    if len(my_list2) >= 1:
                        _name = my_list2[0]
                        _value = item.replace(_name + ":").strip()
                        DLLSunny.SetRequestHeader(self.MessageId, create_string_buffer(_name.encode("utf-8")),
                                                  create_string_buffer(_value.encode("utf-8")))
                return True
            elif isinstance(nweHeaders, dict):
                for key in nweHeaders.keys():
                    DLLSunny.SetRequestHeader(self.MessageId, create_string_buffer(key.encode("utf-8")),
                                              create_string_buffer(nweHeaders[key].encode("utf-8")))
                return True
            else:
                return False
        elif key:
            DLLSunny.SetRequestHeader(self.MessageId, create_string_buffer(key.encode("utf-8")),
                                      create_string_buffer(value.encode("utf-8")))
            return True
        else:
            return False

    def endSend(self):
        '''
        终止发送：使用本命令后,这个请求将不会被发送出去
        :return:
        '''
        DLLSunny.SetResponseHeader(self.MessageId, create_string_buffer("Connection".encode("utf-8")),
                                   create_string_buffer("Close".encode("utf-8")))

    def getHeaders(self, key: str = None):
        '''
        取协议头
        :param key: 不传此参数为取全部协议头
        :return:
        '''
        if key:
            return PointerToText(
                DLLSunny.GetRequestHeader(self.MessageId, create_string_buffer(key.encode("utf-8"))))
        else:
            return PointerToText(DLLSunny.GetRequestAllHeader(self.MessageId))

    def modifyCookies(self, key: str = None, value: str = None):
        '''
        设置或修改cookies
        :param key: 设置 key 值相关 cookie 如 key 为空 设置整个cookies
        :param value: 此值为空时删除 key 值的Cookie 如 key 为空 则删除全部
        :return:
        '''
        if key:
            DLLSunny.SetRequestCookie(self.MessageId, create_string_buffer(key.encode("utf-8")),
                                      create_string_buffer(value.encode("utf-8")))
        else:
            DLLSunny.SetRequestAllCookie(self.MessageId, create_string_buffer(value.encode("utf-8")))

    def getCookies(self, key: str = None, onlyValue: bool = False):
        '''
        获取cookies
        :param key: 传些参数时 只取单条cookie
        :param onlyValue: 是否只取值，key值有效时此参数有效
        :return:
        '''
        if key:
            cookie = PointerToText(DLLSunny.GetRequestCookie(self.MessageId, create_string_buffer(key.encode("utf-8"))))
            if onlyValue == False:
                return cookie
            else:
                return cookie.replace(key + "=", "").replace(";", "").strip()
        else:
            return PointerToText(DLLSunny.GetRequestALLCookie(self.MessageId))

    def getBodyLenght(self):
        '''
        获取Body长度
        :return:
        '''
        return PtrToInt(DLLSunny.GetRequestBodyLen(self.MessageId))

    def getBody(self, toString: bool = False):
        '''
        获取 body
        :param toString: 是否转换为字符串
        :return: bytearray|str
        '''

        '''
        获取Body
        :return: 
        '''
        ptr = DLLSunny.GetRequestBody(self.MessageId)
        body = PtrToByte(ptr, 0, self.getBodyLenght())
        if toString == False:
            return (body)
        else:
            try:
                return (body.decode())
            except:
                return (body.decode("gbk"))


# http 响应操作类
class SunnyResponse:
    def __init__(self, _MessageId):
        self.MessageId = _MessageId

    def getReponseText(self):
        '''
        取响应文本
        :return:
        '''
        return PointerToText(DLLSunny.GetResponseBody(self.MessageId))

    def getHeaders(self, key: str = None):
        '''
        取协议头
        :param key: 不传此参数为取全部协议头
        :return:
        '''
        if key:
            return PointerToText(DLLSunny.GetResponseHeader(self.MessageId, create_string_buffer(key.encode("utf-8"))))
        else:
            return PointerToText(DLLSunny.GetResponseAllHeader(self.MessageId))

    def deleteHeaders(self, key: str = None):
        '''
        删除协议头
        :param key:不传此参删除全部
        :return:
        '''
        if key:
            DLLSunny.DelResponseHeader(self.MessageId, create_string_buffer(key.encode("utf-8")))
        else:
            headers = self.getHeaders()
            header_list = headers.split("\r\n")
            for header in header_list:
                header_kv_list = header.split(":")
                if len(header_kv_list) >= 1:
                    DLLSunny.DelResponseHeader(self.MessageId, create_string_buffer(header_kv_list[0].encode("utf-8")))

    def modifyHeaders(self, nweHeaders: Union[str, dict] = None, key: str = None, value: str = None):
        '''
        新增或修改协议头
        :param nweHeaders: nweHeaders: 可以是 str 或 dict 多条换行
        :param key: 单条 Key
        :param value: 单条 value
        :return: 成功返回真，失败返回假
        '''
        if nweHeaders:
            if isinstance(nweHeaders, str):
                DLLSunny.SetResponseAllHeader(self.MessageId, create_string_buffer(nweHeaders.encode("utf-8")))
                return True
            elif isinstance(nweHeaders, dict):
                nweHeader_str = ""
                for key in nweHeaders.keys():
                    if nweHeader_str:
                        nweHeader_str = nweHeader_str + "\r\n" + key + ":" + nweHeaders[key]
                    else:
                        nweHeader_str = key + ":" + nweHeaders[key]
                    DLLSunny.SetResponseAllHeader(self.MessageId, create_string_buffer(nweHeader_str.encode("utf-8")))
                return True
            else:
                return False
        elif key:
            DLLSunny.SetResponseHeader(self.MessageId, create_string_buffer(key.encode("utf-8")),
                                       create_string_buffer(value.encode("utf-8")))
            return True
        else:
            return False

    def setStatusCode(self, Status: int):
        '''
        修改状态码
        :param Status:
        :return:
        '''
        if Status <= 0:
            DLLSunny.SetResponseStatus(self.MessageId, 200)
        else:
            DLLSunny.SetResponseStatus(self.MessageId, Status)

    def getStatusCode(self):
        '''
        获取状态码
        :return:
        '''
        return PtrToInt(DLLSunny.GetResponseStatusCode(self.MessageId))

    def getStatusCodeText(self):
        '''
        获取状态码对应文本
        :return:
        '''
        res = PointerToText(DLLSunny.GetResponseStatus(self.MessageId))
        a = res.split(" ")
        print(res)
        if len(a) > 1:
            del a[0]
            return " ".join(a).strip()
        return ""

    def getResponseLenght(self):
        '''
        取正文长度
        :return:
        '''
        return PtrToInt(DLLSunny.GetResponseBodyLen(self.MessageId))

    def getResponse(self):
        '''
        取响应体
        :return:
        '''
        return PtrToByte(DLLSunny.GetResponseBody(self.MessageId), 0, self.getResponseLenght())

    def modifyReponse(self, data: Union[str, bytes]):
        if isinstance(data, bytes):
            DLLSunny.SetResponseData(self.MessageId, data, len(data))
        else:
            try:
                data = data.encode("utf-8")
            except:
                data = data.encode("gbk")
            DLLSunny.SetResponseData(self.MessageId, data, len(data))


# http 操作类
class Sunny:
    Request = None
    """ HTTP/HTTPS 请求操作对象 """
    Response = None
    """ HTTP/HTTPS 响应操作对象 """

    def __init__(self, MessageId):
        self.MessageId = MessageId

    def GetRequest(self):
        return (SunnyRequest(self.MessageId))

    def GetResponse(self):
        return (SunnyResponse(self.MessageId))

    def GetRefererIp(self):
        '''
        HTTP/HTTPS 获取发起请求的客户端IP
        :return:
        '''
        return (PointerToText(DLLSunny.GetRequestClientIp(self.MessageId)))
